import tkinter as tk
from tkinter import ttk, messagebox, simpledialog
from tkinter import filedialog
import yaml
import os
import json

class YAMLParamEditor:
    def __init__(self, root):
        self.root = root
        self.root.title("训练参数编辑器")
        self.root.geometry("850x500")
        
        # 配置文件路径（用于保存用户设置的默认路径）
        self.config_file = "yaml_editor_config.json"
        
        # 加载用户设置的默认路径，如果没有则使用初始默认路径
        self.default_path = self.load_default_path() or "E:/ultralytics/cfg/train.yaml"
        
        # 主框架
        self.main_frame = ttk.Frame(root)
        self.main_frame.pack(fill=tk.BOTH, expand=True, padx=10, pady=10)
        
        # 顶部标题
        title_label = ttk.Label(
            self.main_frame, 
            text="训练参数编辑器",
            justify=tk.CENTER, font=("Arial", 10, "bold")
        )
        title_label.pack(pady=10)
        
        # 文件路径显示与选择区域
        self.path_frame = ttk.Frame(self.main_frame)
        self.path_frame.pack(fill=tk.X, padx=10, pady=5)
        
        # 路径显示框
        self.path_var = tk.StringVar()
        path_entry = ttk.Entry(
            self.path_frame, 
            textvariable=self.path_var, 
            state="readonly",
            width=60
        )
        path_entry.pack(side=tk.LEFT, padx=5)
        
        # 路径操作按钮区域
        path_buttons = ttk.Frame(self.path_frame)
        path_buttons.pack(side=tk.LEFT)
        
        # 选择文件按钮
        select_btn = ttk.Button(
            path_buttons, 
            text="选择yaml文件", 
            command=self.select_yaml_file
        )
        select_btn.pack(side=tk.LEFT, padx=2)
        
        # 新增：修改默认路径按钮
        set_default_btn = ttk.Button(
            path_buttons, 
            text="设为默认路径", 
            command=self.set_as_default_path
        )
        set_default_btn.pack(side=tk.LEFT, padx=2)
        
        # 左侧“训练参数”子框架
        self.param_frame = ttk.LabelFrame(self.main_frame, text="训练参数设置")
        self.param_frame.pack(side=tk.LEFT, padx=10, pady=10, fill=tk.Y)
        
        # 定义每个参数的配置
        self.param_config = {
            "Epochs": {
                "type": "number",
                "tooltip": "训练的总轮数，即数据集被完整遍历的次数",
                "options": [10, 20, 50, 100, 200, 300]
            },
            "cache": {
                "type": "boolean",
                "tooltip": "是否启用数据缓存以加速加载",
                "options": [True, False]
            },
            "patience": {
                "type": "number",
                "tooltip": "早停机制耐心值：指标持续不提升的轮数阈值",
                "options": [3, 5, 10, 15, 20]
            },
            "device": {
                "type": "string",
                "tooltip": "训练设备（cpu表示CPU，cuda表示GPU）",
                "options": ["cpu", "cuda", "0", "1", "auto"]
            },
            "batch": {
                "type": "number",
                "tooltip": "每批训练的样本数量",
                "options": [8, 16, 32, 64, 128]
            },
            "workers": {
                "type": "number",
                "tooltip": "数据加载的进程数，通常设为CPU核心数",
                "options": [0, 2, 4, 8, 16]
            },
            "imgsz": {
                "type": "number",
                "tooltip": "输入模型的图像尺寸（如640表示640x640）",
                "options": [320, 480, 640, 800, 1024]
            },
            "project": {
                "type": "string",
                "tooltip": "训练结果保存的根目录",
                "options": ["./runs/train", "./experiments", "./outputs"]
            },
            "save": {
                "type": "boolean",
                "tooltip": "是否保存训练日志、权重等结果",
                "options": [True, False]
            },
            "name": {
                "type": "string",
                "tooltip": "实验名称，用于区分不同训练任务",
                "options": ["exp", "exp1", "exp2", "train", "test"]
            }
        }
        
        self.param_vars = {}  # 存储参数变量
        self.param_entries = {}  # 存储下拉框引用
        
        # 构建参数标签和下拉选择框
        row, col = 0, 0
        for param in self.param_config:
            # 参数标签（带悬停提示）
            param_label = ttk.Label(self.param_frame, text=param, font=("Arial", 9))
            param_label.grid(row=row, column=col*3, padx=5, pady=8, sticky=tk.W)
            param_label.bind("<Enter>", lambda e, p=param: self.show_tooltip(p))
            param_label.bind("<Leave>", lambda e: self.hide_tooltip())
            
            # 参数值变量
            var = tk.StringVar()
            self.param_vars[param] = var
            
            # 根据参数类型设置下拉框状态
            state = "normal" if self.param_config[param]["type"] == "number" else "readonly"
            
            # 创建下拉框
            entry = ttk.Combobox(
                self.param_frame,
                textvariable=var,
                values=self.param_config[param]["options"],
                width=10,
                font=("Arial", 9),
                state=state
            )
            entry.grid(row=row, column=col*3 + 1, padx=5, pady=8)
            
            # 为数字类型添加输入验证
            if self.param_config[param]["type"] == "number":
                validate_cmd = (self.root.register(self.validate_number), '%P')
                entry.configure(
                    validate="key",
                    validatecommand=validate_cmd
                )
            
            self.param_entries[param] = entry
            
            # 每2个参数换一行
            col += 1
            if col % 2 == 0:
                row += 1
                col = 0
        
        # 右侧操作区域
        self.right_frame = ttk.LabelFrame(self.main_frame, text="操作区")
        self.right_frame.pack(side=tk.RIGHT, padx=10, pady=10, fill=tk.BOTH, expand=True)
        
        # 按钮区域
        button_frame = ttk.Frame(self.right_frame)
        button_frame.pack(expand=True)
        
        # 读取文件按钮
        read_btn = ttk.Button(
            button_frame, 
            text="重新加载文件", 
            command=self.read_selected_yaml,
            width=20
        )
        read_btn.pack(pady=10)
        
        # 保存修改按钮
        save_btn = ttk.Button(
            button_frame, 
            text="保存参数到文件", 
            command=self.save_parameters,
            width=20,
            
        )
        save_btn.pack(pady=10)
        
        # 重置按钮
        reset_btn = ttk.Button(
            button_frame, 
            text="重置为当前文件值", 
            command=self.reset_parameters,
            width=20
        )
        reset_btn.pack(pady=10)
        
        # 新增：查看/修改默认路径按钮
        view_default_btn = ttk.Button(
            button_frame, 
            text="查看/修改默认路径", 
            command=self.view_or_edit_default_path,
            width=20
        )
        view_default_btn.pack(pady=10)
        
        # 悬停提示窗口
        self.tooltip = tk.Toplevel(root)
        self.tooltip.withdraw()
        self.tooltip.overrideredirect(True)
        self.tooltip_label = ttk.Label(
            self.tooltip, 
            text="", 
            background="#FFF8DC",
            relief=tk.SOLID, 
            borderwidth=1,
            font=("Arial", 8)
        )
        self.tooltip_label.pack(padx=5, pady=2)
        
        # 存储当前文件的原始数据
        self.current_yaml_data = None
        
        # 自定义按钮样式
        self.style = ttk.Style()
        self.style.configure("Accent.TButton", foreground="white", background="#0066CC", font=("Arial", 10, "bold"))
        
        # 程序启动时自动加载默认路径的文件
        self.auto_load_default_file()

    def validate_number(self, value):
        """验证输入是否为有效数字（空值或正整数）"""
        if not value:  # 允许空值
            return True
        if value.isdigit() and int(value) >= 0:  # 只允许非负整数
            return True
        return False

    def load_default_path(self):
        """从配置文件加载用户设置的默认路径"""
        try:
            if os.path.exists(self.config_file):
                with open(self.config_file, "r", encoding="utf-8") as f:
                    config = json.load(f)
                    return config.get("default_path")
            return None
        except Exception as e:
            print(f"加载配置文件失败: {e}")
            return None

    def save_default_path(self, path):
        """将默认路径保存到配置文件"""
        try:
            config = {"default_path": path}
            with open(self.config_file, "w", encoding="utf-8") as f:
                json.dump(config, f, ensure_ascii=False, indent=2)
            return True
        except Exception as e:
            messagebox.showerror("错误", f"保存默认路径失败: {str(e)}")
            return False

    def set_as_default_path(self):
        """将当前选中的文件路径设为默认路径"""
        current_path = self.path_var.get()
        if not current_path or not os.path.exists(current_path):
            messagebox.showwarning("提示", "请先选择一个有效的yaml文件")
            return
        
        # 确认设置
        confirm = messagebox.askyesno(
            "确认设置", 
            f"确定要将以下路径设为默认路径吗？\n{current_path}\n下次启动程序将自动加载此文件"
        )
        if confirm:
            if self.save_default_path(current_path):
                self.default_path = current_path
                messagebox.showinfo("成功", "默认路径已更新，下次启动将生效")

    def view_or_edit_default_path(self):
        """查看当前默认路径，允许手动编辑"""
        # 显示当前默认路径
        current_default = self.default_path
        new_path = simpledialog.askstring(
            "默认路径设置", 
            f"当前默认路径:\n{current_default}\n\n请输入新的默认路径（留空则保持当前）:",
            parent=self.root
        )
        
        if new_path is None:  # 用户取消
            return
            
        new_path = new_path.strip()
        if not new_path:  # 留空，保持当前
            return
            
        # 验证路径有效性
        if os.path.exists(new_path) and new_path.endswith(".yaml"):
            confirm = messagebox.askyesno(
                "确认修改", 
                f"确定要将默认路径修改为:\n{new_path}?"
            )
            if confirm:
                if self.save_default_path(new_path):
                    self.default_path = new_path
                    messagebox.showinfo("成功", "默认路径已更新，下次启动将生效")
        else:
            messagebox.showerror("错误", "输入的路径无效或不是.yaml文件")

    def auto_load_default_file(self):
        """程序启动时自动加载默认路径的yaml文件"""
        if os.path.exists(self.default_path):
            self.path_var.set(self.default_path)
            self.read_selected_yaml(show_success=False)
        else:
            messagebox.showwarning(
                "提示", 
                f"默认路径的文件不存在：\n{self.default_path}\n请手动选择yaml文件"
            )

    def select_yaml_file(self):
        """手动选择yaml文件"""
        file_path = filedialog.askopenfilename(
            title="选择yaml 文件",
            filetypes=[("YAML Files", "*.yaml"), ("All Files", "*.*")],
            initialdir="."
        )
        if file_path:
            self.path_var.set(file_path)
            self.read_selected_yaml()

    def show_tooltip(self, param):
        """显示参数提示"""
        self.tooltip_label.config(text=self.param_config[param]["tooltip"])
        x, y = self.root.winfo_pointerxy()
        self.tooltip.geometry(f"+{x+12}+{y+12}")
        self.tooltip.deiconify()

    def hide_tooltip(self):
        """隐藏参数提示"""
        self.tooltip.withdraw()

    def read_selected_yaml(self, show_success=True):
        """读取选中的yaml文件"""
        yaml_path = self.path_var.get()
        if not yaml_path:
            messagebox.showwarning("提示", "请先选择yaml文件！")
            return
        
        try:
            with open(yaml_path, "r", encoding="utf-8") as f:
                self.current_yaml_data = yaml.safe_load(f)
            
            # 填充参数到下拉框
            for param in self.param_config:
                if param in self.current_yaml_data:
                    self.param_vars[param].set(str(self.current_yaml_data[param]))
                else:
                    self.param_vars[param].set("")
            
            if show_success:
                messagebox.showinfo("成功", f"已成功加载文件：\n{yaml_path}")
        
        except FileNotFoundError:
            messagebox.showerror("错误", "选中的文件不存在，请重新选择！")
        except yaml.YAMLError:
            messagebox.showerror("错误", "YAML文件格式错误，无法解析！")
        except Exception as e:
            messagebox.showerror("错误", f"加载失败：{str(e)}")

    def reset_parameters(self):
        """重置参数为当前文件中的值"""
        if self.current_yaml_data is None:
            messagebox.showwarning("提示", "请先加载文件！")
            return
        
        # 重新填充参数
        for param in self.param_config:
            if param in self.current_yaml_data:
                self.param_vars[param].set(str(self.current_yaml_data[param]))
            else:
                self.param_vars[param].set("")
        
        messagebox.showinfo("提示", "已重置为当前文件中的参数值")

    def save_parameters(self):
        """保存修改后的参数到YAML文件"""
        yaml_path = self.path_var.get()
        if not yaml_path:
            messagebox.showwarning("提示", "请先选择yaml文件！")
            return
        
        if not os.path.exists(yaml_path):
            messagebox.showerror("错误", "文件不存在，无法保存！")
            return
        
        # 确认是否保存
        confirm = messagebox.askyesno("确认保存", f"确定要将修改保存到以下文件吗？\n{yaml_path}")
        if not confirm:
            return
        
        try:
            # 创建要保存的数据字典
            updated_data = self.current_yaml_data.copy() if self.current_yaml_data else {}
            
            # 从下拉框获取并转换参数值
            for param in self.param_config:
                value_str = self.param_vars[param].get().strip()
                
                # 跳过空值（不修改原文件中的值）
                if not value_str:
                    continue
                
                # 根据参数类型进行转换
                if self.param_config[param]["type"] == "number":
                    # 数字类型参数
                    try:
                        updated_data[param] = int(value_str)
                    except ValueError:
                        messagebox.showerror("错误", f"{param}必须是整数！")
                        return
                elif self.param_config[param]["type"] == "boolean":
                    # 布尔类型参数
                    lower_val = value_str.lower()
                    if lower_val in ["true", "1"]:
                        updated_data[param] = True
                    elif lower_val in ["false", "0"]:
                        updated_data[param] = False
                    else:
                        messagebox.showerror("错误", f"{param}必须是True或False！")
                        return
                else:
                    # 字符串类型参数
                    updated_data[param] = value_str
            
            # 保存到文件
            with open(yaml_path, "w", encoding="utf-8") as f:
                yaml.dump(updated_data, f, sort_keys=False, allow_unicode=True)
            
            # 更新当前数据缓存
            self.current_yaml_data = updated_data
            
            messagebox.showinfo("成功", f"参数已成功保存到：\n{yaml_path}")
        
        except Exception as e:
            messagebox.showerror("错误", f"保存失败：{str(e)}")

if __name__ == "__main__":
    root = tk.Tk()
    app = YAMLParamEditor(root)
    root.mainloop()
    